#include <iostream>
#include <vector>
#include <fstream>
#include <dirent.h>
// #include <opencv2/opencv.hpp>
#include <chrono>
#include "inference.h"
using namespace std;
using namespace cv;
int main()
{
    char lang[] = "LANG=zh_CN.UTF-8";
    char language[] = "LANGUAGE=zh_CN:zh";
    putenv(lang);
    putenv(language);
    setlocale(LC_ALL,"zh_CN.UTF-8");
    string project_base_dir = "/home/ljb/workspace/cpp_project/yolo_deployment/cpp_yolo8/";
    string model_path = project_base_dir+"best.onnx";
    // string model_path = "/home/ljb/temp/yolov5s.onnx";
    string dataset_dir = project_base_dir+ "yolo8_img/";
    vector<string> image_names;
    DIR* dir;
    struct dirent* ent;
    if ((dir = opendir(dataset_dir.c_str()))!=NULL)
    {
        while ((ent = readdir(dir))!=NULL)
        {
            if(ent ->d_type == DT_REG)
            {
                image_names.push_back(dataset_dir+ent->d_name);
            }
        }
    }
    bool using_cuda = false;

    YOLO8 yolo(model_path,Size(640,640),"classes.txt",using_cuda);
    cout<<"yolo"<<endl;
    for(string img_path: image_names)
    {
        Mat frame = imread(img_path);
        // cout<<frame.rows<<endl;
        // cout<<frame.cols<<endl;
        auto start = chrono::steady_clock::now();
        vector<Object> output = yolo.infer(frame);
        auto end = chrono::steady_clock::now();
        auto duration = chrono::duration_cast<chrono::milliseconds>(end-start);
        int objects = output.size();
        cout<<"Number of objects: " <<objects<<endl;
        for (auto obj :output)
        {
            rectangle(frame,obj.box,obj.color,2);
            string pred_string = obj.class_name + ' ' + to_string(obj.confidence).substr(0,4);
            cout<<pred_string<<endl;
            putText(frame,pred_string,Point(obj.box.x+5,obj.box.y-10),FONT_HERSHEY_DUPLEX,1,obj.color);
            putText(frame,"FPS: "+to_string(1000.0/duration.count()).substr(0,4),Point(20,frame.rows-20),FONT_HERSHEY_DUPLEX,0.5,Scalar(0,255,0));
        }
        float scale = 0.8;
        resize(frame,frame,Size(frame.cols*scale,frame.rows*scale));
        imshow("result",frame);
        waitKey(0);

    }
    return 0;
}